package cn.edu.sdjzu.knoverse.controller;

import cn.dev33.satoken.annotation.SaCheckLogin;
import cn.edu.sdjzu.knoverse.model.vo.result.Result;
import cn.edu.sdjzu.knoverse.model.vo.result.ResultCodeEnum;
import cn.edu.sdjzu.knoverse.service.RagService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

import cn.edu.sdjzu.knoverse.service.CustomStreamingChatResponseHandler;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.UUID;
import java.time.Instant;

@Slf4j
@RestController
@RequestMapping("/rag")
@SaCheckLogin
public class RagController {

    @Resource
    private RagService ragService;
    
    private final ExecutorService executorService = Executors.newCachedThreadPool();
    
    private final ObjectMapper objectMapper = new ObjectMapper();

    
    /**
     * RAG增强的流式聊天接口 - 先检索相关内容，再流式输出
     * @param message 用户消息
     * @param retrievalType 检索方式（可选）：VECTOR（向量数据库）、GRAPH（图数据库）、ENHANCED（图增强向量检索）、PARALLEL（并行检索）或AUTO（自动决定）
     * @return SSE响应
     */
    @GetMapping(value = "/chat/stream/rag", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter ragChatStream(
            @RequestParam String message,
            @RequestParam(required = false, defaultValue = "AUTO") String retrievalType) {
        // 创建SSE发射器，设置超时时间为60秒
        SseEmitter emitter = new SseEmitter(60000L);
        
        // 设置超时和完成回调
        emitter.onTimeout(() -> {
            try {
                emitter.send(SseEmitter.event()
                        .name("error")
                        .data("连接超时"));
                emitter.complete();
            } catch (IOException ignored) {
                // 忽略发送超时消息时的异常
            }
        });
        
        emitter.onCompletion(() -> {
            log.info("SSE连接已完成");
        });
        
        // 发送连接建立消息
        try {
            emitter.send(SseEmitter.event().name("open").data("RAG流式连接已建立"));
        } catch (IOException e) {
            emitter.completeWithError(e);
            return emitter;
        }
        
        // 在单独的线程中处理流式回复
        executorService.execute(() -> {
            try {
                // 使用RAG增强的流式模型处理请求
                ragService.streamRagChat(message, retrievalType, new CustomStreamingChatResponseHandler() {
                    @Override
                    public void onPartialResponse(String partialResponse) {
                        try {
                            // 使用INFO级别打印每个token的内容，确保在控制台显示
                            log.info("流式响应token内容: [{}]", partialResponse);
                            
                            // 构建与智谧API类似的JSON响应格式
                            Map<String, Object> responseMap = new HashMap<>();
                            String id = UUID.randomUUID().toString().replace("-", "").substring(0, 20);
                            responseMap.put("id", id);
                            responseMap.put("created", Instant.now().getEpochSecond());
                            responseMap.put("model", "rag-streaming");
                            
                            List<Map<String, Object>> choices = new ArrayList<>();
                            Map<String, Object> choice = new HashMap<>();
                            choice.put("index", 0);
                            
                            Map<String, Object> delta = new HashMap<>();
                            delta.put("role", "assistant");
                            delta.put("content", partialResponse);
                            
                            choice.put("delta", delta);
                            choices.add(choice);
                            
                            responseMap.put("choices", choices);
                            
                            // 转换为JSON字符串
                            String jsonResponse = objectMapper.writeValueAsString(responseMap);
                            
                            // 发送JSON响应
                            emitter.send(SseEmitter.event()
                                    .name("token")
                                    .data(jsonResponse));
                        } catch (IOException e) {
                            onError(e);
                        }
                    }

                    @Override
                    public void onCompleteResponse(ChatResponse completeResponse) {
                        try {
                            // 发送完成事件
                            emitter.send(SseEmitter.event()
                                    .name("complete")
                                    .data("RAG流式回复已完成"));
                            emitter.complete();
                        } catch (IOException e) {
                            onError(e);
                        }
                    }

                    @Override
                    public void onError(Throwable throwable) {
                        try {
                            // 发送错误事件
                            String errorMessage = throwable.getMessage();
                            if (errorMessage == null) {
                                errorMessage = "未知错误";
                            }
                            emitter.send(SseEmitter.event()
                                    .name("error")
                                    .data(errorMessage));
                            emitter.complete();
                        } catch (IOException e) {
                            // 如果发送错误事件失败，直接完成
                            emitter.complete();
                        }
                    }
                    
                    @Override
                    public void onStatus(String type, String status) {
                        try {
                            // 发送状态事件
                            emitter.send(SseEmitter.event()
                                    .name("status")
                                    .data(status));
                        } catch (IOException e) {
                            onError(e);
                        }
                    }
                });
            } catch (Exception e) {
                try {
                    // 发送错误事件
                    String errorMessage = e.getMessage();
                    if (errorMessage == null) {
                        errorMessage = "未知错误";
                    }
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(errorMessage));
                } catch (IOException ioException) {
                    // 忽略发送错误消息时的异常
                } finally {
                    emitter.complete();
                }
            }
        });
        
        return emitter;
    }
    
    /**
     * RAG对话接口（带来源）
     */
    @PostMapping("/chat/with-sources")
    public Result<Map<String, Object>> chatWithSources(@RequestParam String message) {
        dev.langchain4j.service.Result<String> result = ragService.chatWithSources(message);
        
        // 处理来源信息
        List<Map<String, Object>> sources = new ArrayList<>();
        if (result.sources() != null) {
            sources = result.sources().stream()
                .map(content -> {
                    Map<String, Object> sourceMap = new HashMap<>();
                    sourceMap.put("text", content.textSegment().text());
                    
                    // 获取元数据信息（如文件名、时间等）
                    try {
                        // 尝试获取来源文件名
                        if (content.metadata() != null && content.metadata().containsKey("source")) {
                            sourceMap.put("source", content.metadata().get("source"));
                        }
                    } catch (Exception e) {
                        sourceMap.put("source", "未知来源");
                    }
                    
                    return sourceMap;
                })
                .collect(Collectors.toList());
        }
        
        Map<String, Object> resultMap = new HashMap<>();
        resultMap.put("message", result.content());
        resultMap.put("sources", sources);
        
        return Result.success(resultMap);
    }

    /**
     * 获取支持的检索类型和描述
     * @return 检索类型和描述列表
     */
    @GetMapping("/retrieval-types")
    public Result<List<Map<String, String>>> getRetrievalTypes() {
        List<Map<String, String>> retrievalTypes = new ArrayList<>();
        
        Map<String, String> vector = new HashMap<>();
        vector.put("type", "VECTOR");
        vector.put("name", "仅向量数据库");
        vector.put("description", "包含文档内容和详细资料，适合查询具体知识点的详细内容、解释、案例文档等");
        retrievalTypes.add(vector);
        
        Map<String, String> graph = new HashMap<>();
        graph.put("type", "GRAPH");
        graph.put("name", "仅图数据库");
        graph.put("description", "包含知识图谱，适合查询实体间关系，如课程包含哪些知识点、知识点间的前置关系、案例与知识点的关联等");
        retrievalTypes.add(graph);
        
        Map<String, String> enhanced = new HashMap<>();
        enhanced.put("type", "ENHANCED");
        enhanced.put("name", "图增强向量检索");
        enhanced.put("description", "先从图数据库获取相关实体和关系信息，然后利用这些信息增强向量检索，综合两种数据源的优势");
        retrievalTypes.add(enhanced);
        
        Map<String, String> parallel = new HashMap<>();
        parallel.put("type", "PARALLEL");
        parallel.put("name", "并行检索");
        parallel.put("description", "同时使用图数据库和向量数据库，各自独立检索后合并结果");
        retrievalTypes.add(parallel);
        
        Map<String, String> auto = new HashMap<>();
        auto.put("type", "AUTO");
        auto.put("name", "自动");
        auto.put("description", "系统会根据您的问题自动选择最合适的检索方式");
        retrievalTypes.add(auto);
        
        return Result.success(retrievalTypes);
    }
} 